package com.hrbxlh.virtualization.util;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.HashSet;
import java.util.Set;
import java.util.concurrent.TimeUnit;

/**
 * @author cheer
 */
public class LockUtil {

    private static Logger logger = LoggerFactory.getLogger(LockUtil.class);

    private static final Set<String> lockPool = new HashSet<>();

    private static ThreadLocal<Boolean> canUnlock = new ThreadLocal<>();

    public static void lock(String key) {
        while (!addIfNotExists(key)) {
            try {
                Thread.sleep(1000);
            } catch (InterruptedException e) {
                logger.error(e.getMessage(), e);
            }
        }
    }

    public static boolean tryLock(String key) {
        return addIfNotExists(key);
    }

    public static boolean tryLock(String key, long timeout, TimeUnit timeUnit) {
        long millTimeout = timeUnit.toMillis(timeout);
        long startTime = System.currentTimeMillis();
        while (System.currentTimeMillis() - startTime < millTimeout) {
            try {
                if (addIfNotExists(key)) {
                    return true;
                }
                Thread.sleep(millTimeout > 1000 ? 1000 : millTimeout);
            } catch (InterruptedException e) {
                logger.error(e.getMessage(), e);
            }
        }
        return false;
    }

    public static void unlock(String key) {
        remove(key);
    }

    private static synchronized boolean addIfNotExists(String key) {
        if (lockPool.contains(key)) {
            return false;
        }
        lockPool.add(key);
        canUnlock.set(true);
        return true;
    }

    private static synchronized void remove(String key) {
        if (canUnlock.get() != null && canUnlock.get()) {
            lockPool.remove(key);
            canUnlock.remove();
        }
    }
}
